from . import api
from flask import jsonify
from app.model import History
from app import db, scheduler, cache
import numpy as np
from torch import nn
import torch

BATCH_SIZE, TIME_SLOT = 2, 3
score_matrix = None
is_trained = False


# 训练数据
def sequence(batch_size, time_slot):
    from application import app
    with app.app_context():
        # unique = False
        records = torch.tensor(list(map(lambda history: history.pid, History.query.order_by('uid').order_by('date').all())))
    print(records)
    batch = (len(records) - 1) // (batch_size * time_slot) 
    x = records[:batch * time_slot * batch_size].reshape(batch_size, -1)
    y = records[1:batch * time_slot * batch_size + 1].reshape(batch_size, -1)
    for i in range(batch):
        yield x[:, i * time_slot: (i + 1) * time_slot], y[:, i * time_slot: (i + 1) * time_slot]

class RNNLM(nn.Module):
    def __init__(self, batch_size, time_slot, psize, input_size, hidden_size):
        super(RNNLM, self).__init__()
        self.batch_size = batch_size
        self.time_slot = time_slot
        # 先暂定论文数为1024
        self.embedding = nn.Embedding(psize, input_size)
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.linear = nn.Linear(hidden_size, psize)
    
    def forward(self, x):
        o, (self.h, self.c) = self.lstm(self.embedding(x))
        out = self.linear(o)
        return out.contiguous().view(-1, out.size(-1))

model = RNNLM(2, 3, 1024, 3, 3)

def predict(inputs, step):
    inputs = torch.tensor(inputs).unsqueeze(0)
    model.eval()
    result = []
    for i in range(step):
        index = model(inputs)[-2:-1].argmax(dim=1, keepdim=True)
        torch.cat((inputs, index), dim=1)
        result.append(index.item())
    return result

def train_model():
    global is_trained
    # loop over the dataset multiple times
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
    criterion = nn.CrossEntropyLoss()
    for epoch in range(10):
        running_loss = 0.0
        for data in sequence(BATCH_SIZE, TIME_SLOT):
            inputs, labels = data
    
            # zero the parameter gradients
            optimizer.zero_grad()
    
            # forward + backward + optimize
            outputs = model(inputs)
            loss = criterion(outputs, labels.contiguous().view(-1))
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print('Loss: {}'.format(running_loss))
    print('Finished Training')
    is_trained = True
    torch.save(model.state_dict(), './lstm.pt')




def compute_score():
    global score_matrix
    print('compute')
    from application import app
    with app.app_context():
        records = History.query.all()
        umax, pmax = list(db.session.execute("select max(uid), max(pid) from history"))[0]
        item_matrix = np.zeros((umax + 1, pmax + 1), dtype=np.float64)
        for history in records:
            item_matrix[history.uid][history.pid] += 1
        mean_matrix = item_matrix.mean(axis=0)
        correct_matrix = item_matrix.T - mean_matrix.reshape(mean_matrix.size, 1)
        member = np.matmul(correct_matrix, correct_matrix.T)
        std = np.std(correct_matrix, axis=1)
        score_matrix = np.zeros((pmax + 1, pmax + 1))
        for i in range(pmax + 1):
            for j in range(pmax + 1):
                score_matrix[i, j] = member[i, j] / (std[i] * std[j] * umax + 1e-6)

@scheduler.task('interval', id='compute_cross_filter', minutes=10)
def job():
    compute_score()

@scheduler.task('interval', id='lstm', minutes=20)
def train():
    train_model()
    



@api.route('/cross/<int:uid>', methods=['GET'])
@cache.cached(timeout=5 * 60, key_prefix='view:%s')
def cross_filter(uid):
    global score_matrix
    if score_matrix is None:
        compute_score()
    recents = db.session.execute(f"select pid from history where uid = {uid} order by date desc limit 4")
    res = set()
    for recent in recents:
        pid = recent[0]
        index = score_matrix[pid].argsort()[-4: -1]
        res.update(index.tolist())
    return jsonify(code=0, data=list(res))


@api.route('/lstm/<int:uid>', methods=['GET'])
@cache.cached(timeout=5 * 60, key_prefix='view:%s')
def lstm(uid):
    model.load_state_dict(torch.load('./lstm.pt'))
    pids = list(map(lambda history: history.pid , History.query.filter_by(uid=uid).order_by('date').all()))
    print(pids)
    if not pids:
        return jsonify(code=0, data=[])
    result = []
    if model is not None:
        result = predict(pids, 8)
    return jsonify(code=0, data=result)


@api.route('/lstm/train', methods=['GET'])
def request_train():
    train_model()
    return jsonify(code=0, msg='ok')




